use super::types::*;
use super::util::*;

use num_bigint::BigUint;
use num_traits::identities::{One, Zero};

use std::cmp::min;

lazy_static! {
    static ref ONE: Prob = One::one();
}

/***********
 * ENCODER *
 ***********/

/// Supplementary information used by finalize_encode, after the main action is over
#[derive(Debug)]
pub struct EncoderFinalizerState {
    augmented_input: Prob,
    augmented_input_precision: u32,
    // lower_acceptable: BinaryFraction,
    // upper_acceptable: BinaryFraction,
    acceptable_interval_size: BinaryFraction,
}

impl EncoderFinalizerState {
    pub fn new(state: &EncoderState) -> Self {
        let input = &state.remaining_input;
        let f = state.input_significant_bits;

        // left shift whole remaining input to zero
        let remaining_input_shl = shl(input.clone(), f).trunc_as_biguint();

        // input must be encoded fully: lower bound
        let lower_acceptable = shr(BinaryFraction::from_biguint(remaining_input_shl), f);

        let e: BigUint = One::one();

        // upper bound: truncate to the right; don't destroy any input digits! put a 1 *behind*
        let right_truncated_plus_one: BigUint =
            std::ops::Shl::shl(shl(input.clone(), f).trunc_as_biguint(), 1) + e;

        // upper bound (exclusive, value must always stay below this)
        let upper_acceptable: Prob = min(
            shr(
                BinaryFraction::from_biguint(right_truncated_plus_one),
                f + 1,
            ),
            state.lower_bound.clone() + state.probability.clone(),
        );
        Self {
            augmented_input: state.remaining_input.clone(),
            augmented_input_precision: state.input_significant_bits,
            // lower_acceptable: lower_acceptable.clone(),
            // upper_acceptable: upper_acceptable.clone(),
            acceptable_interval_size: upper_acceptable - lower_acceptable,
        }
    }
}

#[derive(Debug)]
pub struct EncoderState {
    accumulated_output: Vec<u8>,
    remaining_input: Prob,
    input_significant_bits: u32,

    lower_bound: Prob,
    probability: Prob,

    processed_input_bits: u32,

    remaining_sentential_form: Vec<Symbol>,

    input_finished: bool,
    finalizer_state: Option<EncoderFinalizerState>,
}

impl EncoderState {
    pub fn new(input_bytes: &[u8]) -> EncoderState {
        let input = bytes_to_binaryfraction_between_0_and_1(input_bytes);
        EncoderState {
            accumulated_output: vec![],
            remaining_input: input.clone(),
            // input_significant_bits: input.denom().trailing_zeros().unwrap() as u32,
            // trailing_zeros of denom 10 is 1, fits
            input_significant_bits: input.f as u32,
            lower_bound: Zero::zero(),
            probability: One::one(),

            processed_input_bits: 0,

            remaining_sentential_form: vec![Symbol::Nonterminal(START_SYMBOL)],

            input_finished: false,
            finalizer_state: None,
        }
    }
}

fn finished_encoding(f: u32, lower: &Prob, prob: &Prob, input: &Prob) -> bool {
    // If the input to be encoded has that many significant bits left, then
    // there's no need for the more expensive check. Beware of random refill!
    if f > 8 {
        return false;
    }

    let upper = lower.clone() + prob;

    let cond1 = shl(input.clone(), f + 1).trunc() - shl(lower.clone(), f + 1).trunc() <= One::one();
    let cond2 = binaryfractions_same_prefix(f, &upper, input);

    cond1 && cond2
}

#[derive(Debug)]
pub enum EncodeStepResult {
    ToBeContinued {
        output: Vec<u8>,
        remaining_input: Prob,
        lower_bound: Prob,
    },
    Final(Vec<u8>),
    PreliminaryFinal,
    Step,
}

/// When (almost) all the input has been encoded, we still need to finalize,
/// for two reasons. One, the sentential form reached may still contain nonterminals
/// and these must be randomly resolved until only terminals are left;
/// two, upper and lower bound may fail to exactly represent the input taken as k/2^{-f}
/// and instead of k/2^{-f} we must encode some k/2^{-f} + d where 0 < d < 2^{-f}.
/// Then d must be a random number uniformly from \[ input, input + 2^{-f} \).
pub fn finalize_encode<'a>(
    grammar: &SCFG,
    state: &mut EncoderState,
    random_source: &'a mut dyn Iterator<Item = u64>,
) -> Result<EncodeStepResult, anyhow::Error> {
    if state.remaining_sentential_form.is_empty() {
        let accumulated_output = state.accumulated_output.clone();
        state.accumulated_output = vec![];

        let upper_bound = state.lower_bound.clone() + state.probability.clone();

        let insufficient_assured_digits = !binaryfractions_same_prefix(
            state.input_significant_bits,
            &state.lower_bound,
            &upper_bound,
        );

        return if insufficient_assured_digits {
            state.remaining_sentential_form = vec![Symbol::Nonterminal(START_SYMBOL)];

            Ok(EncodeStepResult::ToBeContinued {
                output: accumulated_output,
                remaining_input: state.remaining_input.clone(),
                lower_bound: state.lower_bound.clone(),
            })
        } else {
            Ok(EncodeStepResult::Final(accumulated_output))
        };
    }

    let head = state.remaining_sentential_form.remove(0);

    match head {
        Symbol::Terminals(mut w) => {
            state.accumulated_output.append(&mut w);
            Ok(EncodeStepResult::Step)
        }
        Symbol::Nonterminal(n) => {
            if let Some(rule_list) = grammar.rules.get(&n) {
                let g = state
                    .finalizer_state
                    .as_ref()
                    .unwrap()
                    .augmented_input_precision;

                // more or less arbitrary but should be enough - check?
                if g <= 64 {
                    let augmented_input = state
                        .finalizer_state
                        .as_ref()
                        .unwrap()
                        .augmented_input
                        .clone();
                    let acceptable_interval_size = state
                        .finalizer_state
                        .as_ref()
                        .unwrap()
                        .acceptable_interval_size
                        .clone();

                    let r_u64 = random_source.next().unwrap();
                    let r_add_fraction = mk_prob(BigUint::from(r_u64), 64);
                    let r_add_scaled = r_add_fraction * acceptable_interval_size;
                    let r_add_shifted = shr(r_add_scaled, g);

                    state
                        .finalizer_state
                        .as_mut()
                        .unwrap()
                        .augmented_input_precision = g + 64;
                    state.finalizer_state.as_mut().unwrap().augmented_input =
                        augmented_input + r_add_shifted;
                }

                // potentially updated above
                let augmented_input = state
                    .finalizer_state
                    .as_ref()
                    .unwrap()
                    .augmented_input
                    .clone();

                let mut iter = rule_list.iter();

                let rule = loop {
                    match iter.next() {
                        Some(rule) => {
                            let mark = state.lower_bound.clone()
                                + state.probability.clone()
                                    * (rule.cumul_lower.clone() + rule.probability.clone());

                            if mark <= augmented_input {
                                continue;
                            } else {
                                break rule;
                            }
                        }
                        None => {
                            // lax vs. strict treatment of inconsistencies?
                            // break rule_list.last().unwrap();
                            return Err(anyhow::Error::msg(format!(
                                "internal error at nonterminal {}",
                                n
                            )));
                        }
                    }
                };

                let new_lower_bound_increment =
                    state.probability.clone() * rule.cumul_lower.clone();
                let new_lower_bound = state.lower_bound.clone() + new_lower_bound_increment;

                state.lower_bound = new_lower_bound;
                state.probability = state.probability.clone() * rule.probability.clone();

                let mut newx = rule.rhs.clone();
                newx.append(&mut state.remaining_sentential_form[..].to_vec());
                state.remaining_sentential_form = newx;

                Ok(EncodeStepResult::Step)
            } else {
                Err(anyhow::Error::msg(format!(
                    "missing rules for {} in grammar",
                    n
                )))
            }
        }
    }
}

/// Caution: modifies state!
pub fn step_encode(
    grammar: &SCFG,
    state: &mut EncoderState,
) -> Result<EncodeStepResult, anyhow::Error> {
    if shl8(state.lower_bound.clone()).trunc()
        == shl8(state.lower_bound.clone() + state.probability.clone()).trunc()
        && state.input_significant_bits >= 8
    {
        state.processed_input_bits += 8;
        state.input_significant_bits -= 8;
        state.lower_bound =
            shl8(state.lower_bound.clone()) - shl8(state.lower_bound.clone()).trunc();
        state.probability =
            shl8(state.probability.clone()) - shl8(state.probability.clone()).trunc();
        state.remaining_input =
            shl8(state.remaining_input.clone()) - shl8(state.remaining_input.clone()).trunc();
    }

    if finished_encoding(
        state.input_significant_bits,
        &state.lower_bound,
        &state.probability,
        &state.remaining_input,
    ) {
        state.input_finished = true;
        state.finalizer_state = Some(EncoderFinalizerState::new(state));

        let upper = state.lower_bound.clone() + state.probability.clone();
        let really_the_same =
            binaryfractions_same_prefix(state.input_significant_bits, &state.lower_bound, &upper);

        return if state.remaining_sentential_form.is_empty() && really_the_same {
            let accumulated_output = state.accumulated_output.clone();
            state.accumulated_output = vec![];
            Ok(EncodeStepResult::Final(accumulated_output))
        } else {
            Ok(EncodeStepResult::PreliminaryFinal)
        };
    }

    // A word (terminal form) has just been output. This does not mean we're done.
    if state.remaining_sentential_form.is_empty() {
        return if state.remaining_input.is_zero() {
            let accumulated_output = state.accumulated_output.clone();
            state.accumulated_output = vec![];
            state.input_finished = true;
            Ok(EncodeStepResult::Final(accumulated_output))
        } else {
            let accumulated_output = state.accumulated_output.clone();
            state.accumulated_output = vec![];
            // Re-initialize for the next block
            state.remaining_sentential_form = vec![Symbol::Nonterminal(START_SYMBOL)];
            Ok(EncodeStepResult::ToBeContinued {
                output: accumulated_output,
                remaining_input: state.remaining_input.clone(),
                lower_bound: state.lower_bound.clone(),
            })
        };
    }

    // Why not store them in reverse order? Also, try sharing subtrees
    let head = state.remaining_sentential_form.remove(0);

    match head {
        Symbol::Nonterminal(n) => {
            if let Some(rule_list) = grammar.rules.get(&n) {
                let mut new_lower_bound: Prob = state.lower_bound.clone();

                for rule in rule_list.iter() {
                    let increment = state.probability.clone() * rule.probability.clone();
                    let new_upper_bound = new_lower_bound.clone() + increment;

                    if state.remaining_input < new_upper_bound {
                        let mut newx = rule.rhs.clone();
                        newx.append(&mut state.remaining_sentential_form[..].to_vec());

                        state.remaining_sentential_form = newx;
                        state.lower_bound = new_lower_bound;
                        state.probability = state.probability.clone() * rule.probability.clone();

                        return Ok(EncodeStepResult::Step);
                    }

                    // upper bound of this rule is lower bound of the next
                    new_lower_bound = new_upper_bound;
                }
                Err(anyhow::Error::msg(format!(
                    "incomplete rules for {} in grammar",
                    n
                )))
            } else {
                Err(anyhow::Error::msg(format!(
                    "missing rules for {} in grammar",
                    n
                )))
            }
        }
        Symbol::Terminals(mut w) => {
            state.accumulated_output.append(&mut w);
            Ok(EncodeStepResult::Step)
        }
    }
}

pub struct Encoder<'a> {
    grammar: &'a SCFG,
    state: EncoderState,
    error: Option<anyhow::Error>,
    /// Extra bits may be required after encoding the message, to finalize
    random_source: Box<dyn Iterator<Item = u64>>,
}

impl<'a> Encoder<'a> {
    pub fn with_live_random(scfg: &'a SCFG, input: &'a [u8]) -> Encoder<'a> {
        Encoder::new(
            scfg,
            input,
            Box::new(Randoms {
                rng: rand::thread_rng(),
            }),
        )
    }

    pub fn new(
        scfg: &'a SCFG,
        input: &'a [u8],
        random_source: Box<dyn Iterator<Item = u64>>,
    ) -> Encoder<'a> {
        Encoder {
            grammar: scfg,
            state: EncoderState::new(input),
            error: None,
            random_source,
        }
    }
}

/// Encode the given input as a sequence of words of the grammar, by first converting it to an arbitrary-precision binary from [0..1).
/// Error states can only be caused by a faulty grammar.
impl<'a> Iterator for Encoder<'a> {
    type Item = Vec<u8>;

    fn next(&mut self) -> Option<Vec<u8>> {
        if self.state.input_finished {
            loop {
                let final_step_result =
                    finalize_encode(self.grammar, &mut self.state, &mut self.random_source);

                match final_step_result {
                    Err(e) => {
                        self.error = Some(e);
                        return None;
                    }
                    Ok(EncodeStepResult::Final(s)) => {
                        return if s.is_empty() { None } else { Some(s) };
                    }
                    Ok(EncodeStepResult::ToBeContinued { output, .. }) => return Some(output),
                    Ok(EncodeStepResult::Step) => continue,
                    Ok(EncodeStepResult::PreliminaryFinal) => {
                        unreachable!();
                    }
                }
            }
        } else {
            loop {
                let step_result = step_encode(self.grammar, &mut self.state);

                match step_result {
                    Err(e) => {
                        self.error = Some(e);
                        return None;
                    }
                    Ok(EncodeStepResult::Step) => continue,
                    Ok(EncodeStepResult::ToBeContinued { output: s, .. }) => return Some(s),

                    Ok(EncodeStepResult::Final(s)) => return Some(s),
                    // going into the self.state.input_finished == true branch
                    Ok(EncodeStepResult::PreliminaryFinal) => return self.next(),
                }
            }
        }
    }
}
